from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import numpy as np
import cv2
import math
from PIL import Image, ImageDraw, ImageFont
import time
font_path = 'SimHei.ttf'
def draw_chinese(img, text, pos, font_size=32, color=(0, 255, 0)):
    """
    在OpenCV的图像上绘制中文
    Args:
        img: OpenCV图像
        text: 要绘制的文本
        pos: 文本的位置
        font_path: 字体文件的路径
        font_size: 字体的大小
        color: 文本的颜色
    Returns:
        带有文本的OpenCV图像
    """
    # 使用PIL的Image类处理图片
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img_pil)
    # 字体的格式
    font = ImageFont.truetype(font_path, font_size)
    # 在图片上添加文字
    draw.text((pos[0], pos[1]-font_size), text, font=font, fill=color)
    # 转换回OpenCV格式
    img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
    return img

# scripts for crop images
def crop_image(img, position):
    def distance(x1,y1,x2,y2):
        return math.sqrt(pow(x1 - x2, 2) + pow(y1 - y2, 2))    
    position = position.tolist()
    for i in range(4):
        for j in range(i+1, 4):
            if(position[i][0] > position[j][0]):
                tmp = position[j]
                position[j] = position[i]
                position[i] = tmp
    if position[0][1] > position[1][1]:
        tmp = position[0]
        position[0] = position[1]
        position[1] = tmp

    if position[2][1] > position[3][1]:
        tmp = position[2]
        position[2] = position[3]
        position[3] = tmp

    x1, y1 = position[0][0], position[0][1]
    x2, y2 = position[2][0], position[2][1]
    x3, y3 = position[3][0], position[3][1]
    x4, y4 = position[1][0], position[1][1]

    corners = np.zeros((4,2), np.float32)
    corners[0] = [x1, y1]
    corners[1] = [x2, y2]
    corners[2] = [x4, y4]
    corners[3] = [x3, y3]

    img_width = distance((x1+x4)/2, (y1+y4)/2, (x2+x3)/2, (y2+y3)/2)
    img_height = distance((x1+x2)/2, (y1+y2)/2, (x4+x3)/2, (y4+y3)/2)

    corners_trans = np.zeros((4,2), np.float32)
    corners_trans[0] = [0, 0]
    corners_trans[1] = [img_width - 1, 0]
    corners_trans[2] = [0, img_height - 1]
    corners_trans[3] = [img_width - 1, img_height - 1]

    transform = cv2.getPerspectiveTransform(corners, corners_trans)
    dst = cv2.warpPerspective(img, transform, (int(img_width), int(img_height)))
    return dst

def order_point(coor):
    arr = np.array(coor).reshape([4, 2])
    sum_ = np.sum(arr, 0)
    centroid = sum_ / arr.shape[0]
    theta = np.arctan2(arr[:, 1] - centroid[1], arr[:, 0] - centroid[0])
    sort_points = arr[np.argsort(theta)]
    sort_points = sort_points.reshape([4, -1])
    if sort_points[0][0] > centroid[0]:
        sort_points = np.concatenate([sort_points[3:], sort_points[:3]])
    sort_points = sort_points.reshape([4, 2]).astype('float32')
    return sort_points

class Duguang:
    def __init__(self) -> None:
        self.table_recognition = pipeline(Tasks.table_recognition, model='damo/cv_dla34_table-structure-recognition_cycle-centernet')
        self.ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_proxylessnas_ocr-detection-db-line-level_damo', model_revision='v1.0.0')  # 轻量化
        self.ocr_recognition = pipeline(Tasks.ocr_recognition, model='damo/cv_convnextTiny_ocr-recognition-document_damo') # 轻量化
    def infer(self, img_path):
        if isinstance(img_path, str):
            image_full = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)
            image_full = cv2.cvtColor(image_full, cv2.COLOR_BGR2RGB)
        else:
            image_full = np.array(img_path, dtype=np.uint8)
            image_full = cv2.cvtColor(image_full, cv2.COLOR_BGR2RGB)
        print("推理")
        a = time.time()
        det_result = self.ocr_detection(image_full)
        b = time.time()
        print("推理时间", b-a)
        det_result = det_result['polygons'] 
        if det_result.shape[0] == 0:
            return image_full
        print(f"坐标区域为四边形的四个顶点的坐标。这是所有的区域坐标和文字: ")
        xs, ys = [], []
        for i in range(det_result.shape[0]):
            pts = order_point(det_result[i])
            image_crop = crop_image(image_full, pts)
            result = self.ocr_recognition(image_crop)
            pts = pts.astype(np.int32)
            cv2.polylines(image_full, [pts], True, (0, 255, 0), 2)
            print(result['text'][0])
            # font_size : 行的高度
            font_size = int(image_crop.shape[0])
            image_full = draw_chinese(image_full, result['text'][0], ((pts[0][0]), pts[0][1]), font_size=font_size)
        return image_full
if __name__=='__main__':
    dg = Duguang()
    # out = dg.infer('1904.02701v1/1904.02701v1_4.png')